import dask.dataframe as dd
import geopandas as gpd
import matplotlib.pyplot as plt
from dask.diagnostics import ProgressBar
from mpl_toolkits.axes_grid1 import make_axes_locatable
from scipy.interpolate import griddata
import numpy as np
import dask.array as da
#from mpl_toolkits.basemap import Basemap

# 读取地图数据
world = gpd.read_file(gpd.datasets.get_path("naturalearth_lowres"))

# 读取数据
data = dd.read_csv("WCPFC_PS_M_Grid1_Temp_SSH_MLT_2010_2019.csv", usecols=["skj_c_una", "skj_c_log", "skj_c_dfad", "skj_c_afad", "Temp_0", "Temp_50", "Temp_100", "Temp_150", "Temp_200", "Temp_250", "Temp_300", "SSH", "MLT","Lon","Lat"],assume_missing=True)

# 计算Lon和Lat的最小值和最大值

lon_min, lon_max = data["Lon"].min().compute(), data["Lon"].max().compute()
lat_min, lat_max = data["Lat"].min().compute(), data["Lat"].max().compute()

print("Lon range:", lon_min, "-", lon_max)
print("Lat range:", lat_min, "-", lat_max)

# 筛选经纬度范围
lon_min, lon_max, lat_min, lat_max = 110, 180, -10, 10
#lon_min, lon_max, lat_min, lat_max = -179, 180, -51, 46
query_str = f"Lon >= {lon_min} and Lon <= {lon_max} and Lat >= {lat_min} and Lat <= {lat_max}"
data = data.query(query_str)

# 获取列名列表
cols = dd.get_dummies(data).columns.tolist()






# 转换为空间数据框
geometry = gpd.points_from_xy(data["Lon"], data["Lat"])
gdf = gpd.GeoDataFrame(data, geometry=geometry)
gdf.crs = "EPSG:4326"



target_list = ["skj_c_una", "skj_c_log", "skj_c_dfad", "skj_c_afad"]
backend_list = ["Temp_0", "Temp_50", "Temp_100", "Temp_150", "Temp_200", "Temp_250", "Temp_300", "SSH", "MLT"]

for target in target_list:
    for backend in backend_list:
        print(target,backend)

        # 查找target所在的列数
        col_num_skj = cols.index(str(target))
        col_num_lon = cols.index("Lon")
        col_num_lat = cols.index("Lat")
        col_num_temp = cols.index(str(backend))

        s = gdf[col_num_skj]
        s = (s - s.min()) / (s.max() - s.min()) * 100  # 将s缩放到0到100之间

        # in all
        with ProgressBar():
            x = np.linspace(gdf[col_num_lon].min(), gdf[col_num_lon].max(), 100)
            y = np.linspace(gdf[col_num_lat].min(), gdf[col_num_lat].max(), 100)
            xx, yy = np.meshgrid(x, y)
            #points = gdf[[col_num_lon, col_num_lat]].values.compute()
            points = da.from_array(gdf[[col_num_lon, col_num_lat]].values)
            #values = gdf[col_num_temp].values.compute()
            values = da.from_array(gdf[col_num_temp].values)
            grid = griddata(points, values, (xx, yy), method="linear")
            fig, ax = plt.subplots(figsize=(10, 6))
            world.plot(ax=ax, color="white", edgecolor="black")
            ###
            #fig, ax = plt.subplots(figsize=(10, 6))
            #scatter = ax.scatter(gdf[col_num_lon], gdf[col_num_lat], s=gdf[col_num_skj], c="black", alpha=0.5)
            scatter = ax.scatter(gdf[col_num_lon], gdf[col_num_lat], s=s, edgecolor="black", facecolor="none", alpha=0.5)
            gdf.plot(ax=ax, column=col_num_temp, markersize=1, cmap="coolwarm", alpha=0.5, edgecolor="black", facecolor="none", legend=False)
            ax.set_xlim([lon_min, lon_max])
            ax.set_ylim([lat_min, lat_max])
            ax.set_xlabel("Longitude")
            ax.set_ylabel("Latitude")
            ax.set_title(f"Bubble plot of {target} and coolwarm of {backend}")

            # 添加图例
            sizes = [25, 50, 75, 100] # 气泡大小范围
            #labels = [str(size) for size in sizes] # 图例标签
            labels = ['10','50','100','200'] # 图例标签
            handles = [plt.scatter([], [], s=size, edgecolors="black", linewidths=1, color="none") for size in sizes] # 图例标记
            legend1 = ax.legend(handles, labels, loc="lower left", title=str(target), scatterpoints=1, frameon=True, labelspacing=1.5, borderpad=1, handletextpad=2, borderaxespad=1)
            ax.add_artist(legend1)

            ###

            im = ax.imshow(grid.T, extent=[x.min(), x.max(), y.min(), y.max()], origin="lower", cmap="coolwarm")
            divider = make_axes_locatable(ax)
            cax = divider.append_axes("right", size="5%", pad=0.05)
            cbar = fig.colorbar(im, cax=cax)

        plt.savefig(f'output_{target}_{backend}.png')
        plt.close()



